from torch_geometric.datasets import OGB_MAG

dataset = OGB_MAG(root='/home/Dyf/code/dataset/data', preprocess='metapath2vec')
# preprocess也可以用TransE等
data = dataset[0]
print(data)
import torch_geometric.transforms as T
# 为异构图增加反向边
data = T.ToUndirected()(data)
print(data)
"""
HeteroData(
  paper={
    x=[736389, 128],
    year=[736389],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389],
  },
  author={ x=[1134649, 128] },
  institution={ x=[8740, 128] },
  field_of_study={ x=[59965, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  (author, writes, paper)={ edge_index=[2, 7145660] },
  (paper, cites, paper)={ edge_index=[2, 5416271] },
  (paper, has_topic, field_of_study)={ edge_index=[2, 7505078] }
)
HeteroData(
  paper={
    x=[736389, 128],
    year=[736389],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389],
  },
  author={ x=[1134649, 128] },
  institution={ x=[8740, 128] },
  field_of_study={ x=[59965, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  (author, writes, paper)={ edge_index=[2, 7145660] },
  (paper, cites, paper)={ edge_index=[2, 10792672] },
  (paper, has_topic, field_of_study)={ edge_index=[2, 7505078] },
  以下3条边是增加的部分
  # (institution, rev_affiliated_with, author)={ edge_index=[2, 1043998] },
  # (paper, rev_writes, author)={ edge_index=[2, 7145660] },
  # (field_of_study, rev_has_topic, paper)={ edge_index=[2, 7505078] }
)
"""
#以下方法用于判断边节点对应关系是否正确。（通过判断两边最大的节点ID是否超限制来确定）
for edge_type in data.edge_types:
    src, _, dst = edge_type
    assert data[edge_type].edge_index[0].max() < data[src].num_nodes
    assert data[edge_type].edge_index[1].max() < data[dst].num_nodes

dataset = OGB_MAG(root='/home/Dyf/code/dataset/data', preprocess='metapath2vec',transform=T.ToUndirected())
data = dataset[0]
print(data)